import pybullet as p
import pybullet_data
import numpy as np
import time
import math

# 初始化PyBullet
physicsClient = p.connect(p.GUI)  # 可视化模式
# physicsClient = p.connect(p.DIRECT)  # 无界面模式（更快）
p.setAdditionalSearchPath(pybullet_data.getDataPath())  # 添加默认数据路径
p.setGravity(0, 0, -9.8)  # 设置重力

# 加载地面和机器人
# planeId = p.loadURDF("plane.urdf")
robotId = p.loadURDF("rm_65_description/urdf/rm_65_rev.urdf", [0, 0, 0], useFixedBase=True)

# p.resetDebugVisualizerCamera(cameraDistance=0.3, cameraYaw=60, cameraPitch=-50,
#                                     cameraTargetPosition=[0.5, -0.05, 0.1])
p.getCameraImage(320,200, renderer=p.ER_BULLET_HARDWARE_OPENGL )#得到图像

# 获取机器人关节信息
numJoints = p.getNumJoints(robotId)
print(f"机器人共有 {numJoints} 个关节")

# # 下面循环主要用来获取机械臂转动关节数量
# numJoints = p.getNumJoints(robotId) # 
# jointIds_c = []
# for j in range (numJoints):
#      p.changeDynamics(p, j, linearDamping=0, angularDamping=0)
#      info_c = p.getJointInfo(p, j)
#      # print(info)
#      jointName_c = info_c[1]
#      jointType_c = info_c[2]
#      if (jointType_c==p.JOINT_PRISMATIC or jointType_c== p.JOINT_REVOLUTE):
#           jointIds_c.append(j)
# print('转动or移动关节: ', jointIds_c)
# for joint_index in range(numJoints):
#        joint_info = p.getJointInfo(p, joint_index)
#        print(f"\
#                [0]关节索引:{joint_info[0]}\n\
#                [1]关节名称:{joint_info[1]}\n\
#                [2]关节类型:{joint_info[2]}\n\
#                [3]此主体的位置状态变量中的第一个位置索引:{joint_info[3]}\n\
#                [4]在这个物体的速度状态变量中的第一个速度索引:{joint_info[4]}\n\
#                [5]保留参数:{joint_info[5]}\n\
#                [6]关节阻尼大小:{joint_info[6]}\n\
#                [7]关节摩擦系数:{joint_info[7]}\n\
#                [8]滑动或旋转关节的位置下限:{joint_info[8]}\n\
#                [9]滑动或旋转关节的位置上限:{joint_info[9]}\n\
#                [10]关节最大力矩:{joint_info[10]}\n\
#                [11]关节最大速度:{joint_info[11]}\n\
#                [12]连杆名称:{joint_info[12]}\n\
#                [13]在当前连杆坐标系中表示的移动或转动的关节轴(忽略JOINT_FIXED固定关节):{joint_info[13]}\n\
#                [14]在父连杆坐标系中表示的关节位置:{joint_info[14]}\n\
#                [15]在父连杆坐标系中表示的关节姿态(四元数x、y、z、w):{joint_info[15]}\n\
#                [16]父连杆的索引,若是base连杆则返回-1:{joint_info[16]}\n\n")

# 确定末端执行器链接
# 对于KUKA iiwa，末端链接是最后一个可动关节
endEffectorLinkIndex = 6  # 通常最后一个关节是末端执行器

# 获取可动关节索引列表
movableJoints = []
for i in range(numJoints):
    jointInfo = p.getJointInfo(robotId, i)
    jointType = jointInfo[2]
    if jointType == p.JOINT_REVOLUTE or jointType == p.JOINT_PRISMATIC:
        movableJoints.append(i)
print(f"可动关节索引: {movableJoints}")

# 设置初始关节角度（弧度）
initialJointPositions = [0, 0, 0, 0, 0, 0, 0]
for i, jointIndex in enumerate(movableJoints):
    p.resetJointState(robotId, jointIndex, initialJointPositions[i])

# 设置仿真参数
timeStep = 1.0 / 240.0  # 仿真步长
p.setTimeStep(timeStep)

# # 定义目标轨迹（圆形轨迹）
# def generate_circular_trajectory(center, radius, height, num_points=100):
#     """生成3D空间中的圆形轨迹"""
#     trajectory = []
#     for i in range(num_points):
#         angle = 2 * math.pi * i / num_points
#         x = center[0] + radius * math.cos(angle)
#         y = center[1] + radius * math.sin(angle)
#         z = center[2] + height
#         trajectory.append((x, y, z))
#     return trajectory

# # 生成目标轨迹
# trajectory = generate_circular_trajectory(center=[0.5, 0, 0.5], radius=0.2, height=0.1)

trajectory = np.load('myend16.npy')

# 逆运动学求解和轨迹跟踪
def follow_trajectory_with_ik(robotId, endEffectorLinkIndex, trajectory):
    """使用逆运动学跟踪轨迹"""
    joints = []
    
    for targetPos in trajectory:
        # 设置目标位置（可选：设置目标方向）
        # targetOrientation = p.getQuaternionFromEuler([0, math.pi/2, 0])  # 末端朝向
        
        # 计算逆运动学解
        jointPoses = p.calculateInverseKinematics(
            robotId,
            endEffectorLinkIndex,
            targetPos[:3],
            targetOrientation=targetPos[3:],
            maxNumIterations=100,
            residualThreshold=1e-5
        )
        joints.append(jointPoses)
        
        # 应用关节角度
        for i, jointIndex in enumerate(movableJoints):
            p.setJointMotorControl2(
                bodyIndex=robotId,
                jointIndex=jointIndex,
                controlMode=p.POSITION_CONTROL,
                targetPosition=jointPoses[i],
                force=500,
                positionGain=0.1
            )
        
        # 绘制目标点
        p.addUserDebugPoints([targetPos[:3]], [[1, 0, 0]], pointSize=5, lifeTime=0.1)
        
        # 执行仿真步
        p.stepSimulation()
        time.sleep(timeStep)
    return joints
# 运行轨迹跟踪
joints = follow_trajectory_with_ik(robotId, endEffectorLinkIndex, trajectory)
joints = np.array(joints)
np.save('realdmp3d_joint.npy', joints)
p.disconnect()